#pragma once

#include <vector>

#include "arithmetic.h"
#include "lexer.h"
#include "location.h"
#include "semantics.h"
#include "token.h"

struct IRBlock;

namespace IR {
    struct Operand;
}  // namespace IR

namespace AST {
    struct ASTNode {
        IRBlock* irBegin = nullptr;
        IRBlock* irEnd = nullptr;
        IR::Operand* irValue = nullptr;
        IRBlock* irTPos = nullptr;
        IRBlock* irFPos = nullptr;

        Location loc;
        virtual void Print(int dep) const = 0;

        template <class T>
        T* As() {
            return dynamic_cast<T*>(this);
        }

      protected:
        ASTNode() {}
    };

    // Expression - arithmetic expression
    struct Expression : public ASTNode {
        Semantics::Info* info = nullptr;

        void Print(int dep) const override;
        void PrintInfo(int dep) const;
    };

    struct AssignExpr : public Expression {
        ASTNode* lhs = nullptr;  // LHS, Expression
        ASTNode* rhs = nullptr;  // RHS, Expression

        void Print(int dep) const override;
    };

    struct CallExpr : public Expression {
        ASTNode* func = nullptr;  // function name, IdentifierExpr
        ASTNode* args = nullptr;  // arguments, Expression

        void Print(int dep) const override;
    };

    struct ArrayExpr : public Expression {
        ASTNode* lhs = nullptr;  // LHS, Expression
        ASTNode* rhs = nullptr;  // RHS, Expression

        void Print(int dep) const override;
    };

    struct OperatorExpr : public Expression {
        ArithOp op;
        ASTNode* lhs = nullptr;  // LHS, Expression
        ASTNode* rhs = nullptr;  // RHS, Expression

        explicit OperatorExpr(Token::BaseToken* tok) : op((int)tok->tag) {}
        explicit OperatorExpr(ArithOp op) : op(op) {}

        void Print(int dep) const override;
    };

    struct TernaryExpr : public Expression {
        ASTNode* cond = nullptr;
        ASTNode* tVal = nullptr;
        ASTNode* fVal = nullptr;

        void Print(int dep) const override;
    };

    struct LiteralExpr : public Expression {
        ArithValue value;

        explicit LiteralExpr(ArithValue value) : value(value) {}
        explicit LiteralExpr(int i) : value(i) {}
        explicit LiteralExpr(float f) : value(f) {}
        explicit LiteralExpr(Token::BaseToken* tok) {
            auto* lit = tok->As<Token::Literal>();
            Utils::Assert(lit, "expect an literal");
            value = lit->value;
        }

        void Print(int dep) const override;
    };

    struct VoidExpr : public Expression {
        void Print(int dep) const override;
    };

    struct IdentifierExpr : public Expression {
        std::string name;

        explicit IdentifierExpr(std::string name) : name(move(name)) {}
        explicit IdentifierExpr(Token::BaseToken* tok) {
            auto* ident = tok->As<Token::Identifier>();
            Utils::Assert(ident, "expect an identifier");
            name = ident->name;
        }

        void Print(int dep) const override;
    };

    struct InitializerExpr : public Expression {
        std::vector<ASTNode*> items;  // items, vector<Expression*>
        std::vector<std::pair<int, ASTNode*>> parsed;

        inline void Add(ASTNode* item) { items.push_back(item); }

        void Print(int dep) const override;
    };

    // CastExpr - cast from float to int or from int to float
    // (exists only after analyzed)
    struct CastExpr : public Expression {
        Type::Basic from;
        Type::Basic to;
        ASTNode* origin;

        void Print(int dep) const override;
    };

    // AddrExpr - relative addressing, visit an array by offset
    // (exists only after analyzed)
    struct AddrExpr : public Expression {
        ASTNode* base = nullptr;    // array name, IdentifierExpr
        ASTNode* offset = nullptr;  // array index, Expression

        void Print(int dep) const override;
    };

    // ArgList - argument list of a function
    // (exists only after analyzed)
    struct ArgList : public ASTNode {
        std::vector<ASTNode*> items;  // arguments, vector<Expression*>

        inline void Add(ASTNode* item) { items.push_back(item); }

        void Print(int dep) const override;
    };

    // Block - include statements inside
    struct Block : public ASTNode {
        std::vector<ASTNode*> items;
        std::vector<Symbol*> local;

        inline void Add(ASTNode* item) { items.push_back(item); }

        void Print(int dep) const override;
    };

    struct VarDecl : public ASTNode {
        std::string name;
        bool isConst = false;
        Type::Basic baseType = Type::Basic::Void;
        std::vector<ASTNode*> shape;  // shape, vector<Expression*>
        ASTNode* init = nullptr;      // init value, Expression (NULLABLE)

        explicit VarDecl(std::string name) : name(move(name)) {}

        void Print(int dep) const override;
    };

    struct FuncDecl : public ASTNode {
        std::string name;
        Type::Basic retType = Type::Basic::Void;
        std::vector<ASTNode*> args;  // arguments, vector<VarDecl*>;
        ASTNode* content = nullptr;  // content: Block

        explicit FuncDecl(std::string name) : name(move(name)) {}

        void Print(int dep) const override;
    };

    struct IfStmt : public ASTNode {
        ASTNode* cond = nullptr;  // condition, Expression
        ASTNode* tVal = nullptr;
        ASTNode* fVal = nullptr;

        void Print(int dep) const override;
    };

    struct WhileStmt : public ASTNode {
        ASTNode* cond = nullptr;  // condition, Expression
        ASTNode* content = nullptr;

        void Print(int dep) const override;
    };

    struct ContinueStmt : public ASTNode {
        void Print(int dep) const override;
    };

    struct BreakStmt : public ASTNode {
        void Print(int dep) const override;
    };

    struct ReturnStmt : public ASTNode {
        ASTNode* ret = nullptr;  // return value, Expression

        void Print(int dep) const override;
    };

}  // namespace AST
